SageMakerノートブックでオープンソースの大規模言語モデルCerebras-GPTを試してみた

こんにちは。CX事業本部Delivery部のAkkyです。

この数週間（GPT-4が公開されてから2週間しか経っていないのですね）ChatGPTに代表されるLLMの勢いがものすごいですが、AIを開発しているのはOpenAIだけではありません。

CerebrasというHPCを開発している企業が自社のクラスタでLLMを学習し、モデルをApache2.0ライセンスで公開したとの記事を読んだので、さっそく試してみました。

Cerebras Systems Releases Seven New GPT Models Trained on CS-2 Wafer-Scale Systems

Cerebras-GPT: A Family of Open, Compute-efficient, Large Language Models

モデルはHugging Faceで公開されています。

https://huggingface.co/cerebras

Cerebras-GPTはChatGPTのような対話用に訓練されたモデルではないので、そのような使い方は苦手です。また翻訳も苦手です。今回はテキスト補完で実験します。

Amazon SageMaker ノートブックインスタンス

Amazon SageMaker ノートブックインスタンスはフルマネージドでJupyter Notebookを使うことができるサービスです。

今回はg4dn.8xlarge(32vCPU,128GB)インスタンスで実験しました。もっと小さなインスタンスでも動くと思いますが、13Bモデルではメモリを60GBくらい使うので、g4dn.4xlargeが(16vCPU,64GB)ではギリギリ足りないかもしれません。 もちろんもっと小さなモデルであれば、より小さなインスタンスでも試すことができます。

なお、ダウンロードする必要のあるモデル自体のサイズも50GB以上あるので、複数のモデルを試してみる際には200GBくらいのEBSを付けておくと快適に実験できます。私は100GBしかつけなかったのでいちいち消す羽目になりました。

実験の準備

カーネルとしてconda_pytorch_p39を選択し、必要なライブラリをインストールします。 Jupyter Notebookでは、先頭に!を付けることでシステムコマンドを実行できます。

!pip install transformers

また、左のファイル一覧からモデル保存するためのhuggingfaceというディレクトリを作っておきます。

Cerebras-GPTの実験

Hugging Faceにあるサンプルをほぼそのまま使いました。

cache_dirでモデルをEBSに保存するよう変更しています。

CUDA(GPU)で動かすときにはdevice=torch.device("cpu")のcpuをcuda:0とします。ただし、NVIDIA Tesla T4はメモリが16GBなので、大きなモデルは載りません。CPUであれば遅いですが13Bモデルも動きます。

from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
import torch

tokenizer = AutoTokenizer.from_pretrained("cerebras/Cerebras-GPT-13B", cache_dir="/home/ec2-user/SageMaker/huggingface")
model = AutoModelForCausalLM.from_pretrained("cerebras/Cerebras-GPT-13B", cache_dir="/home/ec2-user/SageMaker/huggingface")

text = "Generative AI is"

pipe = pipeline("text-generation", model=model, tokenizer=tokenizer, device=torch.device("cpu"))
generated_text = pipe(text, max_length=300, do_sample=False, no_repeat_ngram_size=2)[0]
print(generated_text['generated_text'])

1.3Bモデルで次のような出力が得られました。機械学習に関する文章が生成できています。

13Bモデルでは次のような出力が得られました。CPUなので出力に10分以上かかります。

日本語の文章も一応出せますが、意味不明な文章が生成されてしまいました。

注意点

いろいろな文章で実験をすると、意味不明な文章や不適切と思われるような文章が生成されることがありました。 あくまでこのモデルは研究・実験用に公開されているものですので、実用にはおすすめできません。

おわりに

オープンソースのLLMであるCerebras-GPTを試してみました。手元に高性能なマシンが無くても、手軽に大規模なモデルを実験できるのはクラウドの良い点です。

オープンソースのモデルが利用できると、ファインチューニングを行ってより要求にマッチしたAIを作れるようになりますね。 今後、日本語に対応したモデルが出てきたりすると、応用例が広がりそうです。